package com.ifast.asset.controller;

import com.baomidou.mybatisplus.mapper.EntityWrapper;
import com.baomidou.mybatisplus.mapper.Wrapper;
import com.baomidou.mybatisplus.plugins.Page;
import com.ifast.asset.domain.*;
import com.ifast.asset.service.*;
import com.ifast.asset.service.impl.AssetsServiceImpl;
import com.ifast.asset.util.SSHHelper;
import com.ifast.asset.util.SSHResInfo;
import com.ifast.common.annotation.Log;
import com.ifast.common.base.BaseController;
import com.ifast.common.utils.Result;
import com.jcraft.jsch.JSchException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.*;

import java.util.Arrays;
import java.util.Date;
import java.util.List;

@Controller
@RequestMapping("/ssh")
public class SSHController extends BaseController {

    @Autowired
    RulesService rulesService;
    @Autowired
    AssetsService assetsService;
    @Autowired
    AntivirusService antivirusService;
    @Autowired
    AntivirusdetailService antivirusdetailService;
    @Autowired
    FilenameService filenameService;
    @Autowired
    WebshelldetailsService webshelldetailsService;
    @Autowired
    WebshellService webshellService;

    @Log("添加端口安全规则")
    @GetMapping("/sendCmd")
    @ResponseBody
    public Result<Object> sendCmd(String id , Integer port,Integer type) {

        Assets assets = assetsService.selectById(id);
        if (assets.getUsername()==null||assets.getPassword()==null){
            return Result.ok("该资产账号密码为空！");
        }
        if(type==0){
            try{
                SSHHelper ssh = new SSHHelper(assets.getIp(),22,assets.getUsername(),assets.getPassword());
                String cmd = "iptables -I INPUT -p tcp --dport "+port+" -j DROP";
                ssh.sendCmd(cmd);
                //保存
                RulesDO rulesDO = new RulesDO();
                rulesDO.setIp(assets.getIp());
                rulesDO.setPort(port);
                rulesDO.setUsername(assets.getUsername());
                rulesDO.setPassword(assets.getPassword());
                rulesDO.setType(type);
                rulesDO.setTime(new Date());
                rulesService.insert(rulesDO);
            }catch (Exception e){
                return  Result.fail();
            }
        }else {
            try{
                SSHHelper ssh = new SSHHelper(assets.getIp(),22,assets.getUsername(),assets.getPassword());
                String cmd = "iptables -I INPUT -p tcp --dport "+port+" -j ACCEPT";
                ssh.sendCmd(cmd);
                //保存
                RulesDO rulesDO = new RulesDO();
                rulesDO.setIp(assets.getIp());
                rulesDO.setPort(port);
                rulesDO.setUsername(assets.getUsername());
                rulesDO.setPassword(assets.getPassword());
                rulesDO.setTime(new Date());
                rulesDO.setType(type);
                rulesService.insert(rulesDO);
            }catch (Exception e){
                return  Result.fail();
            }
        }

        return  Result.ok();
    }

    @ResponseBody
    @PostMapping("/list")
    public Result<Object> list(@RequestBody RulesDO rulesDTO){
        Wrapper<RulesDO> wrapper = new EntityWrapper<>(rulesDTO).orderBy("id", false);
        Page<RulesDO> page = rulesService.selectPage(getPage(RulesDO.class,rulesDTO.getPageNum(),rulesDTO.getPageSize()), wrapper);
        return Result.ok(page);
    }

    @PostMapping( "/batchRemove")
    @ResponseBody
    public Result<String>  remove(@RequestBody Integer[] assetsIds){
        rulesService.deleteBatchIds(Arrays.asList(assetsIds));
        return Result.ok();
    }

    //下发find命令 type=0 为病毒查找   type=1为webshell查找
    @Log("添加入侵检测")
    @GetMapping("/sendFind")
    @ResponseBody
    public Result<Object> sendCmd(String id,Integer type,String url) {

        Assets assets = assetsService.selectById(id);
        String name = assets.getIp();
        if (assets.getUsername()==null||assets.getPassword()==null){

            return Result.ok("该资产账号密码为空！");
        }else if (url==null||"".equals(url)){

            return Result.ok("路径不能为空");
        } else{

            if(type==0){
                AntivirusDO antivirusDO = new AntivirusDO();
                try{
                    //保存任务

                    antivirusDO.setTime(new Date());
                    antivirusDO.setName(name);
                    antivirusService.insert(antivirusDO);
                    SSHHelper ssh = new SSHHelper(assets.getIp(),22,assets.getUsername(),assets.getPassword());
                    //获取病毒库
                    Wrapper wrapper = new EntityWrapper();
                    wrapper.eq("type",0);
                    List <FilenameDO> lists = filenameService.selectList(wrapper);
                    for (int i = 0 ; i < lists.size();i++){

                        String cmd = "find "+url+" -name " + lists.get(i).getName();
                        SSHResInfo sshResInfo = ssh.sendCmd(cmd);
                        if (!"".equals(sshResInfo.getOutRes())){

                            //保存
                            AntivirusdetailDO antivirusdetailDO = new AntivirusdetailDO();
                            antivirusdetailDO.setAntivirusid(antivirusDO.getId());
                            antivirusdetailDO.setDescribe("发现了"+lists.get(i).getAntivirusortype()+"病毒");
                            antivirusdetailDO.setFile(lists.get(i).getName());
                            antivirusdetailDO.setUrl(sshResInfo.getOutRes().replace(lists.get(i).getName(),""));
                            antivirusdetailService.insert(antivirusdetailDO);
                        }
                    }

                }catch (Exception e){
                    antivirusService.deleteById(antivirusDO.getId());
                    return  Result.fail();
                }
            }else {
                WebshellDO webshellDO = new WebshellDO();
                try{
                    webshellDO.setTime(new Date());
                    webshellDO.setName(name);
                    webshellService.insert(webshellDO);
                    SSHHelper ssh = new SSHHelper(assets.getIp(),22,assets.getUsername(),assets.getPassword());
                    //获取病毒库
                    Wrapper wrapper = new EntityWrapper();
                    wrapper.eq("type",1);
                    List <FilenameDO> lists = filenameService.selectList(wrapper);
                    for (int i = 0 ; i < lists.size();i++){

                        String cmd = "find "+url+" -name " + lists.get(i).getName();
                        SSHResInfo sshResInfo = ssh.sendCmd(cmd);
                        if (!"".equals(sshResInfo.getOutRes())){

                            //保存
                            WebshelldetailsDO webshelldetailsDO = new WebshelldetailsDO();
                            webshelldetailsDO.setLevel("轻微");
                            webshelldetailsDO.setUrl(sshResInfo.getOutRes());
                            webshelldetailsDO.setType(lists.get(i).getAntivirusortype());
                            webshelldetailsDO.setWebshellid(webshellDO.getId());
                            webshelldetailsService.insert(webshelldetailsDO);
                        }
                    }

                }catch (Exception e){
                    webshellService.deleteById(webshellDO.getId());
                    return  Result.fail();
                }
            }
        }
        return  Result.ok();
    }


}
